#!/usr/local/bin/python
# -*- encoding:utf-8 -*-

import numpy as np;
import scipy.sparse as sp
from scipy import linalg
import numpy.matlib as matlib
import time
import threading
import os
import numpy.linalg as linalg
from multiprocessing import Process, Manager
import sys
import utils
import argparse

'''This file is designed to use multi core to compute matrix,
Consider Lx, Ly may not sparse matrix
'''

'''This File reads the covariace matrix C_xy, the query synonym words file, 
the document synonym words file'''

config = utils.get_config();
alpha = config.getfloat('rmls_plus', 'alpha')
beta = config.getfloat('rmls_plus', 'beta');

converge_threshhold= config.getfloat('rmls_plus', 'converge_threshhold');
dx = config.getint('rmls_plus', 'dx');
dy = config.getint('rmls_plus', 'dy');
d = config.getint('rmls_plus', 'd');
theta = config.getfloat('rmls_plus', 'theta');
psi = config.getfloat('rmls_plus', 'psi');
phi = config.getfloat('rmls_plus', 'phi')
rU_dict = {};
rV_dict = {};

def saveModel(Lx, Ly, batch):
    print 'Saving Model For Batch %s' % batch;
    LxFile = 'Lx_%s.txt' % batch;
    LyFile = 'Ly_%s.txt' % batch;

    utils.save_sparse_csr(LxFile, Lx.tocsr());
    utils.save_sparse_csr(LyFile, Ly.tocsr());
   

def checkConverge(LPx,LPy, Lx, Ly):
    t = np.linalg.norm(Lx-LPx) + np.linalg.norm(Ly-LPy);
    if t < converge_threshhold:
        return (True, t);
    else:
        return (False, t);

def checkConvergeM(LPx, LPy, Lx, Ly):
    t = utils.sparseFroNorm(LPx - Lx) + utils.sparseFroNorm(LPy - Ly);
    if t < converge_threshhold:
        return (True, t);
    else:
        return (False, t);

def computeConvergeM(LP, L):
    return utils.sparseFroNorm(LP -L);

def getRelationGradientLx(u):
    '''Get the \alpha\sum_qL_{xq}'''
    if rU_dict.has_key(u):
        li = rU_dict[u];
        l = len(li);
        if l !=0:
            sum = li[0][1]*Lx[li[0][0]];
        for i in range(1,l):
            sum  = sum + li[i][1]*Lx[li[i][0]];
    else:
        sum = 0;

    return alpha * sum;

def getRelationGradientLy(v):
    '''Get the \beta\sum_qL_{yq}'''
    if rV_dict.has_key(v):
        li = rV_dict[v];

        l = len(li);
        if l != 0:
            sum = li[0][1]*Ly[li[0][0]];
        else:
            sum = 0;
        for i in range(1, l):
            sum = sum + li[i][1]*Ly[li[i][0]];
    else:
        sum = 0;

    return beta * sum;


def RMLS_Plus(C, CT, T=25, batchno='3'):
    con = -1;
    global Lx, Ly;

    #Generate Lx and Ly;
    Lx = matlib.rand(dx, d);
    Ly = matlib.rand(dy, d);
    manager = Manager();
    t = 0
    con = 0;
    Id = matlib.identity(d);
    psiId = psi * Id;
    phiId = phi * Id;
    while True:
        t += 1;
        LPx = Lx.copy();
        LPy = Ly.copy();
        print 'This is %d iteration, last converge is %f' %(t, con);

        print 'Updating Matrix Lx';

        start = time.time();
        W = C.dot(Ly);
        Y = linalg.inv(theta*Ly.transpose().dot(Ly) + psiId);

        Lx = W.dot(Y);
      
        print 'Update Lx cost: ' + str(time.time() -start); 
    
        print 'Updating Matrix Ly';

        start = time.time();        
        N = CT.dot(Lx);
        X = linalg.inv(theta*Lx.transpose().dot(Lx) + phiId);
        Ly = N.dot(X);
    
        print 'Update Ly cost: ' + str(time.time() - start);

        #Check Converge
        flag, con = checkConverge(LPx, LPy, Lx, Ly);
       
        print 'Converge: ', flag, con;
        if flag or t >= T:
            break;
    Lx = sp.lil_matrix(Lx);
    Ly = sp.lil_matrix(Lx);
    utils.l2norm(Lx, d);
    utils.l2norm(Ly, d);
    return (Lx, Ly);
   
def parse(relational):
    '''this fucntion parse the relational data to a symmetric dict'''
    di = {};

    for line in relational:
        li = line.strip().split('\t');
        i = int(li[0]);
        j = int(li[1]);
        s = float(li[2]);
        if di.has_key(i):
            di[i].append((j,s));
        else:
            di[i] = [];
            di[i].append((j,s));
        if di.has_key(j):
            di[j].append((i, s));
        else:
            di[j] = [];
            di[j].append((i, s));

    return di;

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=' RMLS Plus algorithms for leanring to match from click log and synonyms');
    parser.add_argument('--syn', action='store_true', help='whether to train with synonyms');
    parser.add_argument('--runflag', help= 'the runflag of this run');
    args = parser.parse_args();

    if args.syn:
        relationalU = open('LabeledData/query_I', 'r');
        rU_dict = parse(relationalU);
        relationalU.close();
        relationalV = open('LabeledData/documents_I', 'r');
        rV_dict = parse(relationalV);
        relationalV.close();

    batchno = args.runflag;
    print args.syn, batchno;

    try:
        W = utils.load_sparse_csc('W.npz');
        N = W.transpose().tocsc();
        print N.shape;
        Lx, Ly = RMLS_Plus(N, W, 50, batchno); 
        saveModel(Lx, Ly, batchno);
    except KeyboardInterrupt:
        saveModel(Lx, Ly, batchno);
        sys.exit(0);
